from sqlalchemy import create_engine, Column, Integer, String, DateTime, Float, Text, Boolean, ForeignKey, BigInteger, Enum
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from datetime import datetime
import os
from uuid import uuid4
import enum

# 定义枚举类型
class ScenarioEnum(enum.Enum):
    DEVICE_INSPECTION_I = "设备运检I"
    DEVICE_INSPECTION_S = "设备运检S"
    ENVIRONMENT_PATROL = "环境巡视"
    PERSONNEL_SAFETY = "人员安全"

class OriginEnum(enum.Enum):
    CLOUD = "cloud"
    EDGE = "edge"

# 数据库配置
try:
    from config import config
    engine = create_engine(config.database.url, connect_args={"check_same_thread": False})
except ImportError:
    # 如果配置模块不可用，使用默认的SQLite数据库
    import pathlib
    db_path = pathlib.Path(__file__).parent / "edge_cloud.db"
    engine = create_engine(f"sqlite:///{db_path}", connect_args={"check_same_thread": False})

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

class TrainingTask(Base):
    __tablename__ = "training_tasks"
    
    id = Column(Integer, primary_key=True, index=True)
    task_id = Column(String, unique=True, index=True)
    name = Column(String, default="")  # 新增：任务名称
    model_type = Column(String, default="yolov8n")
    task_type = Column(String, default="detection")
    type = Column(String, default="standard_training")  # 新增：训练类型（如 incremental_learning）
    dataset_name = Column(String, default="10_cls_600")
    status = Column(String, default="pending")  # pending, running, completed, failed
    progress = Column(Integer, default=0)  # 新增：进度百分比
    start_time = Column(DateTime)  # 新增：开始时间
    end_time = Column(DateTime)  # 新增：结束时间
    log_path = Column(String)
    model_output_path = Column(String)
    error_message = Column(Text)  # 存储错误信息
    description = Column(Text)  # 新增：任务描述
    tags = Column(Text)  # 标签，JSON格式存储列表，与models中的tag字段一致
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 新增字段：训练方法选择
    training_method = Column(String, default="incremental_training")  # 训练方法：incremental_training, data_collaboration, model_collaboration
    
    # 新增字段：模型选择
    model_name = Column(String)  # 选择的模型名称
    
    # 新增字段：节点选择
    total_nodes = Column(Integer, default=1)  # 总节点数量
    cloud_node_ids = Column(Text)  # 云侧节点ID列表，JSON格式存储
    edge_node_ids = Column(Text)  # 边侧节点ID列表，JSON格式存储
    
    # 新增字段：训练参数
    training_params = Column(Text)  # 训练参数，JSON格式存储
    
    @property
    def tags_list(self):
        """将JSON字符串格式的tags转换为列表"""
        if self.tags:
            try:
                import json
                return json.loads(self.tags)
            except (json.JSONDecodeError, TypeError):
                return None
        return None
    
    @property
    def cloud_node_ids_list(self):
        """将JSON字符串格式的cloud_node_ids转换为列表"""
        if self.cloud_node_ids:
            try:
                import json
                return json.loads(self.cloud_node_ids)
            except (json.JSONDecodeError, TypeError):
                return []
        return []
    
    @property
    def edge_node_ids_list(self):
        """将JSON字符串格式的edge_node_ids转换为列表"""
        if self.edge_node_ids:
            try:
                import json
                return json.loads(self.edge_node_ids)
            except (json.JSONDecodeError, TypeError):
                return []
        return []
    
    @property
    def training_params_dict(self):
        """将JSON字符串格式的training_params转换为字典"""
        if self.training_params:
            try:
                import json
                return json.loads(self.training_params)
            except (json.JSONDecodeError, TypeError):
                return {}
        return {}
    
    # 训练参数
    epochs = Column(Integer, default=100)
    batch_size = Column(Integer, default=16)
    learning_rate = Column(Float, default=0.001)
    
    # 节点部署信息 - 保留用于兼容性
    cloud_nodes = Column(Integer, default=0)  # 云端节点数量
    edge_nodes = Column(Integer, default=0)   # 边缘节点数量
    
    # 评估结果
    final_accuracy = Column(Float)
    accuracy = Column(Float)  # 新增：当前准确率
    final_loss = Column(Float)
    evaluation_results = Column(Text)  # JSON格式存储详细结果

class Model(Base):
    __tablename__ = "models"
    
    id = Column(Integer, primary_key=True, index=True)
    model_id = Column(String, unique=True, index=True)
    model_name = Column(String)
    model_type = Column(String)
    base_model_name = Column(String)
    dataset_name = Column(String)
    weight_file_path = Column(String)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 新增字段 - 基于提供的JSON结构
    architecture = Column(String)  # 模型架构，如 "ResNet-50"
    version = Column(String)  # 模型版本，如 "v2.3.1"
    status = Column(String, default="draft")  # 模型状态，如 "published", "draft", "deprecated"
    size = Column(String)  # 模型大小，如 "98.3 MB"
    framework = Column(String)  # 框架，如 "PyTorch", "TensorFlow"
    last_trained = Column(String)  # 最后训练时间，如 "2024-12-02"
    deployment_type = Column(String, default="both")  # 部署类型，如 "cloud", "edge", "both"
    description = Column(Text)  # 模型描述
    tags = Column(Text)  # 标签，JSON格式存储列表
    deployment_info = Column(Text)  # 部署信息，JSON格式存储，包含cloudNodes和edgeNodes
    
    # 模型性能指标
    precision = Column(Float)
    recall = Column(Float)
    f1_score = Column(Float)
    map_50 = Column(Float)  # mAP@0.5
    map_95 = Column(Float)  # mAP@0.5:0.95
    
    # 训练信息
    training_task_id = Column(String, ForeignKey("training_tasks.task_id"))
    dataset_id = Column(String, ForeignKey("datasets.dataset_id"))
    
    # 关联关系
    training_task = relationship("TrainingTask", backref="models")
    dataset = relationship("Dataset", backref="models")

class Dataset(Base):
    __tablename__ = 'datasets'
    
    # TODO 目前 dataset_id 和 name 需要保持一致，profile文件中也需要注意dataset_id
    id = Column(Integer, primary_key=True, index=True)
    dataset_id = Column(String, unique=True, nullable=False, index=True)  # 数据集唯一标识
    name = Column(String, nullable=False)
    description = Column(Text)
    uploader = Column(String)  # 上传者
    
    # 学习和标注状态
    is_learned = Column(Boolean, default=False)  # 是否已学习
    is_annotated = Column(Boolean, default=False)  # 是否已标注
    
    # 数据信息
    data_size = Column(String)  # 数据大小，如 "2.3GB"
    data_format = Column(String, default="IMAGE")  # IMAGE, SEQUENCE 等
    sample_count = Column(Integer, default=0)  # 样本数量
    upload_time = Column(String)  # 上传时间字符串格式，如 "2024-01-15"
    
    # 时间和更新信息
    time_range = Column(String)  # 时间范围，如 "2023-2024年"
    update_frequency = Column(String)  # 更新频率，如 "每日"、"每周"
    
    # 数据集特性标识
    is_forgetting = Column(Boolean, default=False)  # 是否遗忘
    is_quantization = Column(Boolean, default=False)  # 是否量化
    is_incremental = Column(Boolean, default=False)  # 是否增量
    is_base = Column(Boolean, default=True)  # 是否基础数据集
    
    # 路径和场景信息
    folder_path = Column(String)  # 文件夹路径
    scenario = Column(Enum(ScenarioEnum), nullable=True)  # 应用场景，限制为指定枚举值
    
    # 新增字段
    base_dataset_id = Column(Integer, ForeignKey("datasets.id", ondelete="SET NULL"), nullable=True, index=True)  # 基础数据集ID
    origin = Column(Enum(OriginEnum), nullable=False, default=OriginEnum.CLOUD)  # 数据来源
    
    # 保留原有的一些有用字段
    dataset_type = Column(String, nullable=False)  # detection, classification, etc.
    status = Column(String, default="active")  # active, inactive, processing
    file_count = Column(Integer, default=0)
    total_size = Column(BigInteger, default=0)  # bytes (保留作为内部使用)
    annotated_samples = Column(Integer, default=0)
    
    # 配置文件路径
    config_path = Column(String)  # dataset.yaml 路径
    profile_path = Column(String)  # profile.yaml 路径
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 自引用关系，用于base_dataset_id
    base_dataset = relationship("Dataset", remote_side=[id], backref="derived_datasets")

class Node(Base):
    __tablename__ = "nodes"
    
    id = Column(Integer, primary_key=True, index=True)
    node_id = Column(String, unique=True, index=True)  # 节点ID，如 'cloud-2', 'edge-1'
    name = Column(String)  # 节点名称，如 '云端节点 2'
    node_type = Column(String)  # 节点类型，如 'cloud', 'edge'
    location = Column(String)  # 位置，如 '上海'
    status = Column(String)  # 状态，如 'online', 'offline'
    service_ip = Column(String)  # 服务端口IP，如 '192.168.1.100:8000'
    cpu = Column(String)  # CPU信息，如 '64 核'
    memory = Column(String)  # 内存信息，如 '256GB'
    gpu = Column(String)  # GPU信息，如 'A100 x8'
    bandwidth = Column(String)  # 带宽信息，如 '10Gbps'
    
    # 新增字段
    running_services = Column(Text)  # 运行服务列表，JSON格式存储，可选值：设备运检I,设备运检S,环境巡视,人员安全
    deployed_models = Column(Text)  # 已部署模型列表，JSON格式存储模型ID列表
    collected_sample_count = Column(Integer, default=0)  # 已归集的样本数量
    uncollected_sample_count = Column(Integer, default=0)  # 未归集的样本数量
    cpu_load = Column(Float, default=0.0)  # CPU负载（百分比）
    memory_load = Column(Float, default=0.0)  # 内存负载（百分比）
    network_load = Column(Float, default=0.0)  # 网络负载（百分比）
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class SchedulingStrategy(Base):
    __tablename__ = "scheduling_strategies"
    
    id = Column(Integer, primary_key=True, index=True)
    strategy_id = Column(String, unique=True, index=True)  # 策略ID
    name = Column(String)  # 策略名称
    description = Column(Text)  # 策略描述
    type = Column(String)  # 策略类型，如 'priority_based', 'load_balanced', 'resource_optimized'
    scenario = Column(String)  # 应用场景，如 '人员作业安监', '设备缺陷识别'
    
    # 时间字段
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    last_executed = Column(DateTime)  # 最后执行时间
    next_execution = Column(DateTime)  # 下次执行时间
    
    # 配置信息，JSON格式存储
    config = Column(Text)  # 策略配置，包含priority、maxConcurrentTasks、resourceAllocation等
    
    # 性能指标，JSON格式存储
    metrics = Column(Text)  # 策略性能指标，包含totalExecutions、successRate、avgExecutionTime等
    
    # 状态字段
    status = Column(String, default="draft")  # 策略状态：draft, active, paused, deprecated
    priority = Column(Integer, default=1)  # 策略优先级，数字越小优先级越高
    
    # 资源分配
    max_concurrent_tasks = Column(Integer, default=1)  # 最大并发任务数
    cloud_nodes = Column(Integer, default=0)  # 分配的云端节点数
    edge_nodes = Column(Integer, default=0)  # 分配的边缘节点数

class TestSet(Base):
    """评测集数据库模型"""
    __tablename__ = "test_sets"
    
    id = Column(Integer, primary_key=True, index=True)
    test_set_id = Column(String, unique=True, index=True)  # 评测集唯一标识
    name = Column(String, nullable=False)  # 评测集名称
    description = Column(Text)  # 评测集描述
    tag = Column(String, nullable=False, index=True)  # 标签，如"设备运检I"等
    is_default = Column(Boolean, default=False)  # 是否为默认评测集
    
    # 原始数据集配置，JSON格式存储
    # 格式: [{"dataset_id": "xxx", "dataset_name": "xxx", "sample_ratio": 0.8}, ...]
    origin_datasets = Column(Text, nullable=False)  # 原始数据集列表及其抽取比例
    
    # 统计信息
    total_samples = Column(Integer, default=0)  # 总样本数
    total_size = Column(String)  # 总数据大小
    
    # 时间字段
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    @property
    def origin_datasets_list(self):
        """将JSON字符串格式的origin_datasets转换为列表"""
        if self.origin_datasets:
            try:
                import json
                return json.loads(self.origin_datasets)
            except (json.JSONDecodeError, TypeError):
                return []
        return []

class EvaluationRecord(Base):
    """模型评测记录表"""
    __tablename__ = "evaluation_records"
    
    id = Column(Integer, primary_key=True, index=True)
    eval_id = Column(String, unique=True, index=True)  # 评测唯一ID（eval目录名）
    model_id = Column(String, index=True)
    model_name = Column(String)
    dataset_id = Column(String, index=True)
    dataset_name = Column(String)
    eval_dir = Column(String)           # 评测根目录绝对路径
    preds_dir = Column(String)          # 可视化目录绝对路径（preds/preds2）
    preds_subdir = Column(String)       # 可视化目录名
    used_samples = Column(Integer, default=0)
    metrics = Column(Text)              # JSON 字符串（map_50/map_95/precision/recall/f1）
    params = Column(Text)               # JSON 字符串（评测参数）
    labels_in_eval = Column(Integer, default=0)
    max_label_class_id = Column(Integer)
    nc = Column(Integer)
    device_type = Column(String)       # 设备类型（如 cpu/cuda/cuda:0）
    status = Column(String, default="completed")  # 评测状态（completed/failed 等）
    created_at = Column(DateTime, default=datetime.utcnow)

# 创建数据库表
def create_tables():
    Base.metadata.create_all(bind=engine)


# 获取数据库会话
def get_db():
    db = SessionLocal()
    try:
        yield db
    except Exception:
        db.rollback()
        raise
    finally:
        db.close() 